#include "huzhou_ros.h"
#include "huzhou_common.h"

#include "ServerThread.h"

#include "conf.h"

#include <stdlib.h>
#include <regex>

#include <chrono>
#include <thread>

#include "joyButtonDetector.h"



namespace {

    #define SLEEP_PERIOD_MS    300

    bool audio_buffer_empty = false;

    bool kws_ctrl_running = true;

    enum class SysStatus {
        WAIT_FOR_KWS,        // 正在等待检测关键词
        WAIT_FOR_ASR,        // 正在等待检测到语音输入
        WAIT_FOR_AUDIO_END   // 正在等待大模型结束和语音播放结束
    };

    SysStatus sys_status = SysStatus::WAIT_FOR_KWS;


    // 进行关键词识别的回调函数
    int asr_inputStream_callback(const float* audioData, int numSamples) {
        if (hzc::kws) {
            hzc::kws->DetectKeyword(static_cast<int>(hzc::asr->sample_rate()), audioData, numSamples);
        }
        return 0;
    }


    void waitForTtsAudioEmpty(void) {
        // 等待 tts 文本缓冲区空，也就是语音合成只剩下最后一句了
        while (!hzc::tts->bufferEmpty()) {
            sys_log("waiting for tts buffer empty......");
            std::this_thread::sleep_for(std::chrono::milliseconds(SLEEP_PERIOD_MS));
        }
        sys_log("tts buffer empty!");
    
        // 设置 tts 回调函数
        hzc::tts->setBufferTransitionToEmptyCallback([]() {
            hzc::tts->fillSilence();
            audio_buffer_empty = true;
        });
    
        // 等待语音播放完成
        while (!audio_buffer_empty) {
            sys_log("waiting for audio buffer empty......");
            std::this_thread::sleep_for(std::chrono::milliseconds(SLEEP_PERIOD_MS));
        }

        audio_buffer_empty = false;
    
        hzc::tts->setBufferTransitionToEmptyCallback(nullptr);
        sys_log("audio buffer empty!");
    }


    static ServerThread openAsrThread([](){
        // TTS 处理
        hzc::tts->appendText("你好！我在呢！");
        hzc::tts->notifyInputFinished();
        
        waitForTtsAudioEmpty();

        // 规定rebot说完打招呼的语言之后才是等待asr，如果在rebot打招呼的时候打断会产生bug
        sys_status = SysStatus::WAIT_FOR_ASR;

        // 切换到 ASR
        hzc::asr->setInStreamCallbackProcessAudio();
        hzc::asr->resume();
    });

    // 回调函数，当检测到关键词时调用，这个函数会在 kws->DetectKeyword 中被调用
    void OnKeywordDetected(const std::string &keyword) {
        if (!kws_ctrl_running) {
            sys_log("key words spotting was banned.");
            return;
        }

        // 识别到了关键词
        sys_log("Detected keyword: %s", keyword.c_str());

        sys_log("pause asr audio stream."); // 暂停音频数据输入和识别
        hzc::asr->pause();
        hzc::tts->setMute(false);  // 设置扬声器静音


        // 把后续工作放入线程中处理，保证尽快退出回调函数，防止音频数据积累
        openAsrThread.notify();
    }

    static ServerThread sayHelloThread([](){
        hzc::tts->setMute(false);

        hzc::tts->appendText("大家好，我叫小知，我是湖州吴兴知识产权问答语音智能体。"
            "我可以提供知识产权领域智能问答服务，为您的创新之路保驾护航。");
        hzc::tts->notifyInputFinished();
        
        waitForTtsAudioEmpty();

        sys_status = SysStatus::WAIT_FOR_KWS;

        // 切换到 KWS 识别
        hzc::asr->setReadInStreamCallback(asr_inputStream_callback);
        hzc::asr->resume();  // 恢复语音数据输入

        sys_status = SysStatus::WAIT_FOR_KWS;
    });


    static ServerThread leaveThread([](){
        hzc::tts->setMute(false);
        
        hzc::tts->appendText("好的，期待再次为大家服务，再见。");
        hzc::tts->notifyInputFinished();
        
        waitForTtsAudioEmpty();

        sys_status = SysStatus::WAIT_FOR_KWS;

        // 切换到 KWS 识别
        hzc::asr->setReadInStreamCallback(asr_inputStream_callback);
        hzc::asr->resume();  // 恢复语音数据输入

        sys_status = SysStatus::WAIT_FOR_KWS;
    });


// modify ----------------------------------------->
// ===================================================================================

static std::string asrResStr;


// 预编译正则表达式 - 按照替换目标命名
static const std::regex RE_TO_ROLLER("滚桶|滚筒"); // 替换为"辊筒"
static const std::regex RE_TO_HUZHOU("福州|欧洲"); // 替换为"湖州"
static const std::regex RE_TO_WUXING("吴清|吴京|吴香|卫星|武神|吴欣|无心|五星|无锡|无线|武兴"); // 替换为"吴兴"
static const std::regex RE_TO_WUXING_DISTRICT("吴承旭|无情区"); // 替换为"吴兴区"
static const std::regex RE_TO_TRADEMARK("双倍"); // 替换为"商标"
static const std::regex RE_TO_LOGISTICS("五六"); // 替换为"物流"
static const std::regex RE_TO_PATENT("装地"); // 替换为"专利"
static const std::regex RE_TO_MUQIAN("目前"); // 替换为"专利"




static const std::regex RECEPT_WUXING("吴兴吴兴");
static const std::regex RECEPT_ZHUANLI("专利专利");
static const std::regex RECEPT_BUJU("布局布局");


static bool is_button_Y_pressed = false;

static bool panjie_1_flag = false;

static bool panjie_2_flag = false;

static ServerThread asrSendThread([](){

    // 进行正则表达式的替换
    // 手动修复asr的识别盲区
    asrResStr = std::regex_replace(asrResStr, RE_TO_ROLLER, "辊筒");
    asrResStr = std::regex_replace(asrResStr, RE_TO_HUZHOU, "湖州");
    asrResStr = std::regex_replace(asrResStr, RE_TO_WUXING, "吴兴");
    asrResStr = std::regex_replace(asrResStr, RE_TO_WUXING_DISTRICT, "吴兴区");
    asrResStr = std::regex_replace(asrResStr, RE_TO_TRADEMARK, "商标");
    asrResStr = std::regex_replace(asrResStr, RE_TO_LOGISTICS, "物流");
    asrResStr = std::regex_replace(asrResStr, RE_TO_PATENT, "专利");
    // asrResStr = std::regex_replace(asrResStr, RE_TO_MUQIAN, "当前");


    // 进行重复词语的替换
    asrResStr = std::regex_replace(asrResStr, RECEPT_WUXING, "吴兴");
    asrResStr = std::regex_replace(asrResStr, RECEPT_ZHUANLI, "专利");
    asrResStr = std::regex_replace(asrResStr, RECEPT_BUJU, "布局");

    sys_log("%s", asrResStr.c_str());

    if (panjie_1_flag) {

        panjie_1_flag = false;

        // 第一个问题的套话
        hzc::tts->appendText("我识别到您刚刚说的是，湖州市知识产权发展情况。我正在深入分析这个问题。请您稍等片刻。");
        waitForTtsAudioEmpty();

        sleep(1);

        hzc::tts->appendText("截至2024年底，湖州市知识产权发展成效显著。全市有效发明专利"
            "21579件，同比增长11.55%，高价值发明专利9288件。累计注册商标140621"
            "件，申请中10117件马德里商标有效登记305件。此外，湖州市还拥有45"
            "个地理标志、7个集成电路布图设计，专利密集型产品已备案1812个，"
            "已认定110个，数据知识产权已登记594项，标志着数据要素市场化配置"
            "的规范化进程加速。这些数据充分展示了湖州市在知识产权保护和创新方"
            "面的积极努力和卓越成就。");

        hzc::asr->setReadInStreamCallback(asr_inputStream_callback); // 切换到关键词识别的callback

        hzc::asr->resume();  // 恢复语音输入

        hzros::pubQuestionToUi(asrResStr);

        sys_status = SysStatus::WAIT_FOR_KWS;

        return;

    } else if (panjie_2_flag) {

        panjie_2_flag = false;
       
        // 第二个问题的套话
        hzc::tts->appendText("我识别到您刚刚说的是，物流装备产业海外专利布局情况。我正在深入分析这个问题。请您稍等片刻。");
        waitForTtsAudioEmpty();

        sleep(1);

        hzc::tts->appendText("湖州市物流装备产业在海外专利布局方面主要分为输送分拣设"
                             "备与仓储搬运设备两大领域，截至2024年底，在输送分拣设备领域，湖"
                             "州市共有10件PCT申请;在仓储搬运设备领域，湖州市共有4件PCT申"
                             "请。");

        hzc::asr->setReadInStreamCallback(asr_inputStream_callback); // 切换到关键词识别的callback

        hzc::asr->resume();  // 恢复语音输入

        hzros::pubQuestionToUi(asrResStr);

        sys_status = SysStatus::WAIT_FOR_KWS;

        return;
    }

    // 给区长的问题留一个后门
    else if (is_button_Y_pressed) {

        is_button_Y_pressed = false;

        // 给区长说的话变成套话
        hzc::tts->appendText("我识别到您刚刚说的是，吴兴区目前注册商标。我正在深入分析这个问题。请您稍等片刻。");
        waitForTtsAudioEmpty();

        sleep(1);

        hzc::tts->appendText("截至2024年底，湖州市吴兴区累计注册商标38404件，申请中2042件，马德里商标有效登记58件。"
            "根据湖州市2024年全年新增注册商标的整体趋势，其商标注册量在2025年可能继续保持增长。");

        hzc::asr->setReadInStreamCallback(asr_inputStream_callback); // 切换到关键词识别的callback

        hzc::asr->resume();  // 恢复语音输入

        hzros::pubQuestionToUi(asrResStr);

        sys_status = SysStatus::WAIT_FOR_KWS;

        return;

    }
 
    else {

        // 先发送一些套话，同时等待大模型的回答
        // hzc::tts->appendText("我识别到您刚刚说的是，" + asrResStr + "。我正在深入分析这个问题。请您稍等片刻。");

        hzc::tts->appendText("小知同学听到啦。请您稍等片刻。");

        hzc::tts->notifyInputFinished();

        hzros::pubQuestion(asrResStr); // 向大语言模型节点发送 question
        hzros::pubQuestionToUi(asrResStr);
    }


});


// 语音识别识别到了一些内容，发送给湖州大模型
AsrStream::asrContinueEnum asr_callback(std::string &res) {

    // if (res.length() < 8) {
    // sys_log("ignore words. <<<<<<<<<<<<<<<<<<<<<<<<<");
    // return AsrStream::asrContinue;
    // }

    // ==============================================================
    // 给自我介绍添加一个后门，绕过大模型直接进行回答
    if (res.find("自我介绍") != std::string::npos ||
    res.find("介绍一下你自己") != std::string::npos) {
        hzc::asr->pause();

        sys_log("say hello: %s", res.c_str());

        sayHelloThread.notify();

        return AsrStream::asrContinue;
    }
    // ==============================================================


    // ==============================================================
    // 给退场添加一个后门，绕过大模型直接进行回答
    if (res.find("原路返回") != std::string::npos) {
    hzc::asr->pause();

    sys_log("leave stage: %s", res.c_str());

    leaveThread.notify();

    return AsrStream::asrContinue;
    }
    // ==============================================================

    sys_status = SysStatus::WAIT_FOR_AUDIO_END;

    hzc::asr->pause(); // 暂停音频数据的输入

    // 把剩下的内容全部交给另外一个线程
    asrResStr = res;
    asrSendThread.notify();

    return AsrStream::asrContinue;
}



// ================================================================================

    // 等待语音播放完成，恢复 语音输入 和 关键词检测
    static ServerThread resumeKwsThread([]() {
        sys_log("llm server done, notify input finished.");
        hzc::tts->notifyInputFinished();

        // 等待语音播放完成
        waitForTtsAudioEmpty();

        // hzc::tts->fillSilence(); // ?

        // 恢复关键词识别
        hzc::asr->setReadInStreamCallback(asr_inputStream_callback); // 切换到关键词识别的callback

        hzc::asr->resume();  // 恢复语音数据输入

        sys_status = SysStatus::WAIT_FOR_KWS;

        sys_log("Bingo! you can say keywords again.");
    });

    // 从 湖州大模型 中接受到了一些内容：
    void llmRecStreamData(const std_msgs::String::ConstPtr& msg_p) {
        if (msg_p->data.empty()) {
            return;
        }

        sys_log("rec: %s", msg_p->data.c_str());

        // 判断是否是结束，如果以 "[DONE]" 开头的话，就说明大模型的回答已经结束了
        if (msg_p->data.compare(0, 6, "[DONE]") == 0) {

            // 把后续工作放入一个线程中，尽快退出回调函数，否则音频区会积累音频数据
            resumeKwsThread.notify();

        } else {
            // 添加到语音合成的 buffer 里面，进行语音的合成
            hzc::tts->appendText(msg_p->data);
        }
    }


    void stopTts(void) {
        if (sys_status == SysStatus::WAIT_FOR_KWS) {
            return; // 等关键词的时候不中断，避免bug
        }else if (sys_status == SysStatus::WAIT_FOR_AUDIO_END) {
            audio_buffer_empty = true;
        }
        
        hzc::tts->setMute(true);  // tts 静音， tts 生成的音频数据不会放入缓冲区

        sys_log("received 'stop' msgs......");

        // 保证 tts 语音输出清空
        hzc::tts->clearAudioBuffer();
        hzc::tts->textBufferReset();
        hzc::tts->fillSilence();

        // hzc::tts->manualRunTransitionToEmptyCB();
         
        // 恢复关键词识别
        hzc::asr->setReadInStreamCallback(asr_inputStream_callback); // 切换到关键词识别的callback
        hzc::asr->resume();  // 恢复语音数据输入

        sys_status = SysStatus::WAIT_FOR_KWS;
    }


    // 停止 tts 的函数
    void stop_sig_sub_callback(const std_msgs::String::ConstPtr& msg_p) {
        stopTts();
    }


    // 接收 tts 服务
    void tts_sig_sub_callback(const std_msgs::String::ConstPtr& msg_p) {
        hzc::tts->appendText(msg_p->data);
        hzc::tts->notifyInputFinished();
    }

    void stopKwsServer(void) {
        kws_ctrl_running = false;
        stopTts();
        sys_log("stop kws ctrl running.");
    }

    void resumeKwsServer(void) {
        kws_ctrl_running = true;
        sys_log("resume kws ctrl running.");
    }

    // 控制服务
    void ctrl_sig_sub_callback(const std_msgs::Bool::ConstPtr& msg) {
        if (msg->data) {
            kws_ctrl_running = true;
            sys_log("resume kws ctrl running.");
        } else {
            kws_ctrl_running = false;
            stopTts();
            sys_log("stop kws ctrl running.");
        }

    }


    // 跳过 kws
    void skipKws() {
        if (sys_status != SysStatus::WAIT_FOR_KWS) {
            sys_log("sys status != WAIT_FOR_KWS");
            return;
        }

        sys_log("skip kws.");

        hzc::asr->pause();

        // 取消静音
        hzc::tts->setMute(false);

        // TTS 处理
        hzc::tts->appendText("你好！我在呢！");
        hzc::tts->notifyInputFinished();
        
        waitForTtsAudioEmpty();

        // 规定rebot说完打招呼的语言之后才是等待asr，如果在rebot打招呼的时候打断会产生bug
        sys_status = SysStatus::WAIT_FOR_ASR;

        // 切换到 ASR
        hzc::asr->setInStreamCallbackProcessAudio();
        hzc::asr->resume();
    }


    void quzhangSaySentences(void) {
        is_button_Y_pressed = true;
    }

    // robot 的结束语
    void robotExitPlayAudio(void) {
        hzc::tts->setMute(false);
        hzc::tts->appendText("好的，期待再次为大家服务，再见。");
    }

    // robot 的开场白
    void robotSayHello(void) {
        hzc::tts->setMute(false);
        hzc::tts->appendText("大家好，我叫小知，我是湖州吴兴知识产权问答语音智能体。"
            "我可以提供知识产权领域智能问答服务，为您的创新之路保驾护航。");
    }


    void panjie_1_question_cb(void) {
        panjie_1_flag = true;

    }

    void panjie_2_question_cb(void) {
        panjie_2_flag = true;

    }
}


int main(int argc, char  *argv[])
{
    hzc::init();  // 初始化 asr, tts 等

    // 初始化 ros 相关的数据
    hzros::init(argc, argv, llmRecStreamData, 
        stop_sig_sub_callback, 
        tts_sig_sub_callback,
        ctrl_sig_sub_callback);

    JoyButtonDetector button_detector(skipKws, stopKwsServer, resumeKwsServer,
        quzhangSaySentences,

        // panjie_1_question_cb,
        // panjie_2_question_cb,
        nullptr,
        nullptr,


        robotSayHello,
        robotExitPlayAudio,
        nullptr,
        nullptr);

    // 语音识别的回调函数
    hzc::asr->setAsrCallback(asr_callback);
    // 切换到关键词识别模型
    hzc::asr->setReadInStreamCallback(asr_inputStream_callback);
    // 检测到关键词的 callback func
    hzc::kws->SetKeywordDetectedCallback(OnKeywordDetected);

    // 恢复语音输入
    hzc::asr->resume();

    sys_log("you can say the keywords.");

    hzros::spin();

    hzc::signalHandler(0);
}
